from PIL import Image, ImageDraw, ImageFont
import os
import subprocess
import time
import psutil
import textwrap
import sys
from multiprocessing import Pool, cpu_count

# 预加载题目数据
qlistQuest = []
qlistAns = []
with open('resources/questlist.txt', 'r', encoding='utf-8') as f:
    current_question = []
    for line in f:
        clean_line = line.strip()
        if not clean_line:
            continue
        
        if "正确答案" in clean_line:
            qlistAns.append(clean_line)
            qlistQuest.append(current_question)
            current_question = []
        else:
            wrapped = textwrap.wrap(clean_line, width=25, 
                                   replace_whitespace=False)
            current_question.extend(wrapped)

class PerformanceMonitor:
    def __init__(self):
        self.start_time = time.time()
        self.cpu_percent = 0.0
        
    def update(self):
        self.cpu_percent = psutil.cpu_percent(interval=0.5)
        
    @property
    def elapsed(self):
        return time.time() - self.start_time

def format_time(seconds):
    """将秒数转换为易读的时间格式"""
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return f"{h:02.0f}:{m:02.0f}:{s:05.2f}"

def init_dirs():
    os.makedirs('video', exist_ok=True)
    os.makedirs('image', exist_ok=True)

def generate_images(main_img, font, id):
    try:
        image = main_img.copy()
        draw = ImageDraw.Draw(image)
        y_position = 400
        
        for text in qlistQuest[id]:
            draw.text((100, y_position), text, font=font)
            y_position += 70
        
        question_frame = image.copy()
        answer_text = qlistAns[id]
        draw.text((100, y_position + 200), answer_text, 
                 fill=(255,0,0), font=font)
        answer_frame = image.copy()
        
        return question_frame, answer_frame
    except IndexError:
        print(f"\n⚠️ 题目{id}数据不完整，已跳过")
        return None, None

def create_video(id, q_frame, a_frame):
    if q_frame is None or a_frame is None:
        return False
    
    os.makedirs(f'image/{id}', exist_ok=True)
    
    temp_files = {
        'question': f'image/{id}/temp_q.png',
        'answer': f'image/{id}/temp_a.png'
    }
    
    try:
        q_frame.save(temp_files['question'])
        a_frame.save(temp_files['answer'])
        
        video_path = f'video/{id}.mp4'
        ffmpeg_cmd = [
            'ffmpeg', '-y',
            '-loop', '1', '-t', '4', '-i', temp_files['question'],
            '-loop', '1', '-t', '2', '-i', temp_files['answer'],
            '-filter_complex', '[0][1]concat=n=2:v=1:a=0',
            '-c:v', 'libx264', '-preset', 'ultrafast',
            '-pix_fmt', 'yuv420p', video_path
        ]
        subprocess.run(ffmpeg_cmd, check=True, 
                      stdout=subprocess.DEVNULL, 
                      stderr=subprocess.DEVNULL)
        return True
    except Exception as e:
        print(f"\n⚠️ 视频生成失败（题目{id}）: {str(e)}")
        return False
    finally:
        for f in temp_files.values():
            if os.path.exists(f):
                os.remove(f)

def add_audio(id):
    try:
        video_file = f'video/{id}.mp4'
        final_output = f'video/{id}_final.mp4'
        
        ffmpeg_cmd = [
            'ffmpeg', '-y',
            '-i', video_file,
            '-i', 'resources/music.mp3',
            '-c', 'copy', '-shortest',
            final_output
        ]
        subprocess.run(ffmpeg_cmd, check=True,
                      stdout=subprocess.DEVNULL,
                      stderr=subprocess.DEVNULL)
        
        os.replace(final_output, video_file)
        return True
    except Exception as e:
        print(f"\n⚠️ 音频合并失败（题目{id}）: {str(e)}")
        return False

def process_question(id):
    try:
        main_image = Image.open('resources/main.png')
        font = ImageFont.truetype('resources/font.ttf', 50)
        
        q_img, a_img = generate_images(main_image, font, id)
        if not create_video(id, q_img, a_img):
            return False
        return add_audio(id)
    except Exception as e:
        print(f"\n⚠️ 处理失败（题目{id}）: {str(e)}")
        return False

def print_progress(current, total, speed, cpu_usage, elapsed, remaining):
    """自定义进度显示函数"""
    progress = current / total
    bar_length = 40
    filled_length = int(bar_length * progress)
    bar = '█' * filled_length + '-' * (bar_length - filled_length)
    
    sys.stdout.write(
        f"\r🎬 生成进度 |{bar}| {current}/{total} "
        f"[剩余 {remaining}] "
        f"[速度 {speed:.2f}题/秒] "
        f"[CPU {cpu_usage:.1f}%] "
        f"[耗时 {elapsed}]"
    )
    sys.stdout.flush()

if __name__ == '__main__':
    init_dirs()
    total = len(qlistAns)
    pm = PerformanceMonitor()
    
    print(f"📚 总题目数: {total}")
    print("🚀 开始处理...")
    
    success_count = 0
    last_print_time = 0
    try:
        with Pool(processes=cpu_count()) as pool:
            results = pool.imap_unordered(process_question, range(total))
            
            for i, result in enumerate(results, 1):
                pm.update()
                current_time = time.time()
                
                # 更新成功计数
                if result:
                    success_count += 1
                
                # 计算统计信息
                elapsed = format_time(pm.elapsed)
                speed = i / pm.elapsed if pm.elapsed > 0 else 0
                remaining = format_time((total - i) / speed) if speed > 0 else "--:--:--"
                
                # 每秒更新一次显示
                if current_time - last_print_time >= 0.5 or i == total:
                    print_progress(
                        current=i,
                        total=total,
                        speed=speed,
                        cpu_usage=pm.cpu_percent,
                        elapsed=elapsed,
                        remaining=remaining
                    )
                    last_print_time = current_time
                
            # 完成时换行
            print()
            
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断操作！")
    finally:
        total_time = pm.elapsed
        print("\n" + "="*40)
        print(f"✅ 处理完成！")
        print(f"⏱️ 总耗时: {total_time:.2f}秒")
        print(f"📈 成功率: {success_count}/{total} ({success_count/total:.1%})")
        print(f"🚀 平均速度: {success_count/total_time:.2f}题/秒")
        print(f"🔥 峰值CPU占用: {pm.cpu_percent:.1f}%")